/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.buffer;

import com.aduib.boot.common.concurrent.JctoolsHelper;
import com.aduib.boot.common.log.LOG;

import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @description: DirectBufferPool
 * @author: zzh
 * @date: 2021/12/1 16:17
 */
public class DirectBufferPool implements BufferPool {

  private final int chunkSize;

  private final int pageSize;

  private final short pageCount;

  private final int conReadBuferChunk;

  /** 记录对线程ID->该线程的所使用Direct Buffer的size */
  private final ConcurrentMap<Long, Long> memoryUsage;

  private BufferPage[] allPages;

  private AtomicLong prevAllocatedPage;

  public DirectBufferPool(int pageSize, short chunkSize, short pageCount, int conReadBuferChunk) {
    allPages = new BufferPage[pageCount];
    this.chunkSize = chunkSize;
    this.pageSize = pageSize;
    this.pageCount = pageCount;
    this.conReadBuferChunk = conReadBuferChunk;
    prevAllocatedPage = new AtomicLong(0);
    for (int i = 0; i < pageCount; i++) {
      allPages[i] = new BufferPage(BufferUtils.allocateBuffer(pageSize), chunkSize);
    }
    memoryUsage = JctoolsHelper.nonBlockingHashMap();
  }
  /**
   * TODO 当页不够时，考虑扩展内存池的页的数量...........
   *
   * @param buffer
   * @return
   */
  public Buffer expandBuffer(Buffer buffer) {
    int oldCapacity = buffer.capacity();
    int newCapacity = oldCapacity << 1;
    Buffer newBuffer = allocate(newCapacity);
    if (newBuffer != null) {
      int newPosition = buffer.writerIndex();
      newBuffer.put(buffer);
      newBuffer.position(newPosition);
      recycle(buffer);
      return newBuffer;
    }
    return null;
  }

  @Override
  public Buffer allocate(int size) {
    final int theChunkCount = size / chunkSize + (size % chunkSize == 0 ? 0 : 1);
    int selectedPage = (int) (prevAllocatedPage.incrementAndGet() % allPages.length);
    Buffer byteBuf = allocateBuffer(theChunkCount, 0, selectedPage);
    if (byteBuf == null) {
      byteBuf = allocateBuffer(theChunkCount, selectedPage, allPages.length);
    }
    final long threadId = Thread.currentThread().getId();

    if (byteBuf != null) {

      // 这里必须加锁，因为并发情况下如果allocate和recycle函数操作同一个数据，假设它们都先get到数据，然后allocate先put操作，
      // recycle后进行put操作，这样allocate的put的数据就被覆盖掉
      final Buffer finlBuffer = byteBuf;
      memoryUsage.compute(
          threadId,
          (aLong, aLong2) -> {
            if (aLong2 != null) {
              return memoryUsage.get(threadId) + finlBuffer.capacity();
            } else {
              return (long) finlBuffer.capacity();
            }
          });
    }

    if (byteBuf == null) {
      return BufferUtils.allocateBuffer(size, false);
    }
    return byteBuf;
  }

  @Override
  public PooledBuffer allocate(PoolThreadCache cache, int reqCapacity, int maxCapacity) {
    throw new UnsupportedOperationException();
  }

  @Override
  public BigBuffer allocateBig() {
    return new BigBuffer(this);
  }

  @Override
  public BigBuffer allocateBig(int reqCapacity) {
    return new BigBuffer(this, reqCapacity);
  }

  @Override
  public void recycle(Buffer theBuf) {
    // 堆内buffer直接就清空就好
    if (theBuf != null && (!theBuf.isDirect())) {
      theBuf.clear();
      return;
    }

    final long size = theBuf.capacity();

    boolean recycled = false;
    StringBuilder relatedThreadId = new StringBuilder();

    Buffer thisNavBuf = theBuf; //
    int chunkCount = theBuf.capacity() / chunkSize; // chunk的个数
    Buffer parentBuf = theBuf.attachment(); // page的DirectBuffer
    int startChunk = (int) ((thisNavBuf.address() - parentBuf.address()) / chunkSize); // 开始chunk的序号
    for (int i = 0; i < allPages.length; i++) { // 在所有的页面中查找当前buffer分配的
      if ((recycled
          == allPages[i].recycleBuffer(
              parentBuf, theBuf, startChunk, chunkCount, relatedThreadId))) {
        break;
      }
    }

    final Long threadId =
        relatedThreadId.length() > 0
            ? Long.parseLong(relatedThreadId.toString())
            : Thread.currentThread().getId();

    memoryUsage.computeIfAbsent(threadId, aLong -> (long) (memoryUsage.get(threadId) - size));

    if (recycled == false) {
      LOG.w("warning ,not recycled buffer " + theBuf);
    }
  }

  @Override
  public void recycle(PoolThreadCache cache, Buffer theBuf) {
    throw new UnsupportedOperationException();
  }

  private Buffer allocateBuffer(int theChunkCount, int startPage, int endPage) {
    for (int i = startPage; i < endPage; i++) {
      Buffer buffer = allPages[i].allocatChunk(theChunkCount);
      if (buffer != null) {
        prevAllocatedPage.getAndSet(i);
        return buffer;
      }
    }
    return null;
  }

  @Override
  public int getChunkSize() {
    return chunkSize;
  }

  @Override
  public ConcurrentMap<Long, Long> getNetDirectMemoryUsage() {
    return memoryUsage;
  }

  @Override
  public int getPageSize() {
    return pageSize;
  }

  public short getPageCount() {
    return pageCount;
  }

  @Override
  public long capacity() {
    return (long) pageSize * pageCount;
  }

  @Override
  public long size() {
    return (long) pageSize * chunkSize * pageCount;
  }

  // TODO
  @Override
  public int getSharedOptsCount() {
    return 0;
  }

  public BufferPage[] getAllPages() {
    return allPages;
  }

  @Override
  public int getConReadBufferChunk() {
    return conReadBuferChunk;
  }
}
